from torch.utils.data import DataLoader

from cnn_train.CnnStockConfig import CnnStockConfig
from cnn_train.CnnStockDataProvider import CnnStockDataProvider
from cnn_train.CnnStockDataSet import CnnStockDataSet
from cnn_train.CnnStockModel import CnnStockModel
from cnn_train.CnnStockTrain import CnnStockTrain

if __name__ == '__main__':
    provider = CnnStockDataProvider()
    train_data, train_value, test_data, test_value = provider.get_predicated_data()

    train_dataset = CnnStockDataSet(train_data, train_value)
    test_dataset = CnnStockDataSet(test_data, test_value)

    train_loader = DataLoader(train_dataset, batch_size=CnnStockConfig.batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_dataset, batch_size=CnnStockConfig.batch_size, shuffle=False, drop_last=True)

    # init model
    stockModel = CnnStockModel().to(CnnStockConfig.device)

    # model train
    CnnStockTrain(stockModel, train_loader, test_loader).train()
